/*
 * Description: defender action utilities.
 */
#include "defender_utils.h"
#include <sstream>
#include <vector>
#include "text_util/string_proc.h"

namespace PERF {

std::unordered_map<ExclusionType, std::wstring> DefenderUtils::EXCLUSION_TYPE_TO_COM_PARAM = {
    { ExclusionType::PATH, L"ExclusionPath" },
    { ExclusionType::PROCESS, L"ExclusionProcess" },
    { ExclusionType::EXTENSION, L"ExclusionExtension" },
};
const _bstr_t CLS_NAME_MPPREFERENCE = L"MSFT_MpPreference";

bool DefenderUtils::AddExclusionItem(const std::string &item, ExclusionType type)
{
    bool ret = false;
    do {
        CComPtr<IWbemLocator> pLocator = nullptr;
        CComPtr<IWbemServices> pServices = nullptr;
        CComPtr<IWbemClassObject> pParamInstance = nullptr;
        if (!DefenderCommandInit(pLocator, pServices)) {
            break;
        }
        _bstr_t methodName{};
        if (!DefenderCommandGetMethod(ExclusionCommand::ADD, pServices, pParamInstance, methodName) ||
            !DefenderCommandExec(item, pServices, pParamInstance, methodName, type)) {
            break;
        }
        ret = true;
    } while (0);
    CoUninitialize();
    return ret;
}

bool DefenderUtils::RemoveExclusionItem(const std::string &item, ExclusionType type)
{
    bool ret = false;
    do {
        CComPtr<IWbemLocator> pLocator = nullptr;
        CComPtr<IWbemServices> pServices = nullptr;
        CComPtr<IWbemClassObject> pParamInstance = nullptr;
        _bstr_t methodName{};
        if (!DefenderCommandInit(pLocator, pServices) ||
            !DefenderCommandGetMethod(ExclusionCommand::REMOVE, pServices, pParamInstance, methodName)) {
            break;
        }
        ret = DefenderCommandExec(item, pServices, pParamInstance, methodName, type);
    } while (0);
    CoUninitialize();
    return ret;
}

bool DefenderUtils::GetCurrentExclusionItems(std::vector<std::string> &items, ExclusionType type)
{
    CComPtr<IWbemLocator> pLocator = nullptr;
    CComPtr<IWbemServices> pServices = nullptr;
    if (!DefenderCommandInit(pLocator, pServices)) {
        LOG_ERR("[PERF - GetCurrentExclusionItems] - failed to init");
    }
    auto entry = EXCLUSION_TYPE_TO_COM_PARAM.find(type);
    if (entry == EXCLUSION_TYPE_TO_COM_PARAM.end()) {
        LOG_ERR("[PERF - GetCurrentExclusionItems] - failed, reason: ExclusionType mismatched");
        return false;
    }
    std::wstringstream queryCmdStream;
    _bstr_t queryLang(L"WQL");
    queryCmdStream << L"select " << entry->second << L" from MSFT_MpPreference";
    _bstr_t queryStr(queryCmdStream.str().c_str());
    CComPtr<IEnumWbemClassObject> spQueryResult = nullptr;
    HRESULT res = pServices->ExecQuery(queryLang, queryStr, WBEM_FLAG_RETURN_IMMEDIATELY | WBEM_FLAG_FORWARD_ONLY,
                                       nullptr, &spQueryResult);
    if (spQueryResult == nullptr) {
        LOG_ERR("[PERF - GetExistingExclusionItems] - ExecQuery failed, hresult: {}", res);
        return false;
    }

    ULONG uObjectCount = 0;
    CComPtr<IWbemClassObject> spWmiObject = nullptr;
    spQueryResult->Next(WBEM_INFINITE, 1, &spWmiObject, &uObjectCount);
    if (spWmiObject == nullptr) {
        LOG_ERR("[PERF - GetExistingExclusionItems] - fetch query result failed, hresult: {}", res);
        return false;
    }

    CComVariant cvExclusionPath;
    res = spWmiObject->Get(entry->second.c_str(), 0, &cvExclusionPath, nullptr, nullptr);
    if (FAILED(res)) {
        LOG_ERR("[PERF - GetExistingExclusionItems] - get exclusionPaths failed, hresult: {}", res);
        return false;
    }
    return ReadExclusionItemsFromCComVariant(cvExclusionPath, items);
}

bool DefenderUtils::ReadExclusionItemsFromCComVariant(CComVariant &variant, std::vector<std::string> &items)
{
    auto pArr = variant.parray;
    if (!pArr) {  // exclusionPath list is empty
        return true;
    }
    for (long idx = 0; idx < pArr->rgsabound[0].cElements; idx++) {
        long idxs[1] = { idx };
        BSTR exclusionPathBStr = nullptr;
        SafeArrayGetElement(pArr, idxs, &exclusionPathBStr);
        if (exclusionPathBStr == nullptr) {
            return false;
        }
        LPSTR exclusionPathCStr = _com_util::ConvertBSTRToString(exclusionPathBStr);
        if (exclusionPathCStr == nullptr) {
            SysFreeString(exclusionPathBStr);
            return false;
        }
        items.emplace_back(std::string(exclusionPathCStr));
        SysFreeString(exclusionPathBStr);
        delete[] exclusionPathCStr;
    }
    return true;
}

bool DefenderUtils::DefenderCommandInit(CComPtr<IWbemLocator> &pLocator, CComPtr<IWbemServices> &pServices)
{
    LOG_DEBUG("[PERF] - DefenderExclusionAccelerator defender command init.");
    // Init COM interface
    HRESULT hr = CoInitializeEx(0, COINIT_MULTITHREADED);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator CoInitializeEx init COM failed.");
        return false;
    }

    // Init WMI instance
    hr = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID *)&pLocator);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator init WMI instance failed.");
        return false;
    }

    // Get IWbemService pointer
    if (pLocator == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pLocator is nullptr.");
        return false;
    }
    hr = pLocator->ConnectServer(_bstr_t(L"ROOT\\Microsoft\\Windows\\Defender"), NULL, NULL, 0, NULL, 0, 0, &pServices);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator get IWbenService failed.");
        return false;
    }

    // Set safety parameters
    hr = CoSetProxyBlanket(pServices, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL,
                           RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator safety parameter set failed.");
        return false;
    }
    return true;
}

bool DefenderUtils::DefenderCommandGetMethod(ExclusionCommand cmd, CComPtr<IWbemServices> pServices,
                                             CComPtr<IWbemClassObject> &pParamInstance, _bstr_t &methodName)
{
    HRESULT hr;
    CComPtr<IWbemClassObject> pClass = nullptr;
    CComPtr<IWbemClassObject> pInSignature = nullptr;
    if (pServices == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pServices is nullptr.");
        return false;
    }
    hr = pServices->GetObject(CLS_NAME_MPPREFERENCE, 0, NULL, &pClass, NULL);
    switch (cmd) {
        case ExclusionCommand::ADD:
            methodName = _bstr_t(L"Add");
            break;
        case ExclusionCommand::REMOVE:
            methodName = _bstr_t(L"Remove");
            break;
        default:
            methodName = _bstr_t();
            break;
    }
    if (pClass == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pClass is nullptr.");
        return false;
    }
    hr = pClass->GetMethod(methodName, 0, &pInSignature, NULL);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator MSFT_MpPreference add method get failed.");
        return false;
    }

    if (pInSignature == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pInSignature is nullptr.");
        return false;
    }
    hr = pInSignature->SpawnInstance(0, &pParamInstance);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator spawn instance failed.");
        return false;
    }
    return true;
}

bool DefenderUtils::DefenderCommandExec(const std::string &exclusiveItem, CComPtr<IWbemServices> pServices,
                                        CComPtr<IWbemClassObject> pParamInstance, _bstr_t &methodName,
                                        ExclusionType type)
{
    if (pParamInstance == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pParamInstance is nullptr.");
        return false;
    }
    if (pServices == nullptr) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator pServices is nullptr.");
        return false;
    }
    auto entry = EXCLUSION_TYPE_TO_COM_PARAM.find(type);
    if (entry == EXCLUSION_TYPE_TO_COM_PARAM.end()) {
        LOG_WARN("[PERF - GetCurrentExclusionItems] - failed, reason: ExclusionType mismatched");
        return false;
    }
    HRESULT hr;
    // Create an array
    SAFEARRAYBOUND rgsaBounds[1];
    rgsaBounds[0].cElements = 1;
    rgsaBounds[0].lLbound = 0;
    CComSafeArray<BSTR> psaStrings(rgsaBounds, 1);

    // Add a string to the array
    CComVariant vString;
    V_VT(&vString) = VT_BSTR;
    V_BSTR(&vString) = _bstr_t(exclusiveItem.c_str());
    LONG lArrayIndex = 0;
    SafeArrayPutElement(psaStrings, &lArrayIndex, V_BSTR(&vString));

    // variant array
    CComVariant vStringList;
    V_VT(&vStringList) = VT_ARRAY | VT_BSTR;
    V_ARRAY(&vStringList) = psaStrings;

    // Store the value for the in parameters

    hr = pParamInstance->Put(entry->second.c_str(), 0, &vStringList, CIM_STRING | CIM_FLAG_ARRAY);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator put exclusion path list failed.");
        return false;
    }

    // Execute defender method
    hr = pServices->ExecMethod(CLS_NAME_MPPREFERENCE, methodName, 0, NULL, pParamInstance, NULL, NULL);
    if (FAILED(hr)) {
        LOG_ERR("[PERF] - DefenderExclusionAccelerator execute remove defender failed.");
        return false;
    }

    LOG_INFO("[PERF] - DefenderExclusionAccelerator disapply successfully.");
    return true;
}

}  // namespace PERF